import torch as pt
import numpy as np
import os
import sys

np.random.seed(777)
pt.manual_seed(777)
np.set_printoptions(edgeitems=200)
pt.set_printoptions(edgeitems=200)
VER = 'v2.0'
FILENAME = os.path.basename(__file__)
SAVE_DIR = os.path.join('_save', FILENAME, VER)
os.makedirs(SAVE_DIR, exist_ok=True)
SAVE_PATH = os.path.join(SAVE_DIR, 'weights.dat')
SAVE_PATH4DICT = os.path.join(SAVE_DIR, 'dict.dat')

# 2.	使用pytorch实现循环神经网络处理（每题2分）
# (1)	数据处理
# ①	载入数据，将数据中转换成字典类型进行处理
sentence = """\
This is a long sentence for testing RNN by Pytorch. \
And we will use this paragraph as testing and training material. \
Let's see the power of Pytorch!\
"""
if os.path.exists(SAVE_PATH4DICT):
    dict, idx2char, char2idx = pt.load(SAVE_PATH4DICT)
    print('Loaded dictionary.')
else:
    dict = set(list(sentence))
    idx2char = list(dict)
    char2idx = {ch: i for i, ch in enumerate(idx2char)}
    pt.save((dict, idx2char, char2idx), SAVE_PATH4DICT)
    print('Saved dictionary.')

# ②	设置常用参数，序列长度10，层数为2，循环500次
ALPHA = 0.001
N_STEPS = 10
N_LAYERS = 2
N_ITERS = 500
N_DICT_LEN = len(idx2char)
N_RNN_HIDDEN = 64

# ③	处理x，y用之前的序列长度，预测下个序列长度的结果
x_str = []
y_str = []
x_idx = []
y_idx = []
for i in range(len(sentence) - N_STEPS):
    x_str_i = sentence[i:i+N_STEPS]
    y_str_i = sentence[i+1:i+1+N_STEPS]
    x_str.append(x_str_i)
    y_str.append(y_str_i)
    x_idx.append([char2idx[ch] for ch in x_str_i])
    y_idx.append([char2idx[ch] for ch in y_str_i])

# ④	将y标签变为独热编码形式
x = np.eye(N_DICT_LEN)[x_idx]
y = np.eye(N_DICT_LEN)[y_idx]

# ⑤	将预处理的结果存储成张量形式
x = pt.from_numpy(x).float()
y_idx = np.array(y_idx)
y = pt.from_numpy(y_idx).reshape(-1).long()


# (2)	创建模型类
# ①	创建构造函数，定义相对参数
class MyRnn(pt.nn.Module):

    def __init__(self):
        super().__init__()
        # ②	定义LSTM模型进行处理
        self.rnn = pt.nn.LSTM(N_DICT_LEN, N_RNN_HIDDEN, N_LAYERS, batch_first=True)
        # ③	全连接进行处理
        self.fc = pt.nn.Linear(N_RNN_HIDDEN, N_DICT_LEN)

    # ④	创建正向传播流程
    # ⑤	初始化状态设为0
    def forward(self, x):
        x, _ = self.rnn(x)  # 不传入状态，就是初始化状态设为0
        x = x.reshape(-1, N_RNN_HIDDEN)
        x = self.fc(x)
        return x


# (3)	模型处理
# ①	调用创建好的模型
model = MyRnn()

# ②	配合合理的损失函数和梯度优化算法
criterion = pt.nn.CrossEntropyLoss()
optim = pt.optim.Adam(params=model.parameters(), lr=ALPHA)


def acc(h, y):
    return h.argmax(1).eq(y.long()).float().mean()

# ③	实现梯度下降流程
if os.path.exists(SAVE_PATH):
    print('Loading weights...')
    model.load_state_dict(pt.load(SAVE_PATH))
    print('Loaded')
else:
    for i in range(N_ITERS):
        model.train(True)
        optim.zero_grad()
        h = model(x)
        cost = criterion(h, y)
        cost.backward()
        optim.step()
        model.train(False)
        cost = cost.detach().numpy()
        accv = acc(h, y).detach().numpy()
        # ④	打印每次代价损失值
        print(f'#{i + 1}: cost = {cost}, acc = {accv}')
        pt.save(model.state_dict(), SAVE_PATH)
        print('Saved')

# ⑤	打印预测的结果
print('打印预测的结果')
model.eval()
pred = model(x)
pred = pred.reshape(-1, N_STEPS, N_DICT_LEN).argmax(2)
pred = [''.join([idx2char[i] for i in predi]) for predi in pred]
for x_line, p_line, y_line in zip(x_str, pred, y_str):
    print(f'|{x_line}| => |{p_line}| (hypothesis)')
    print(f'|{x_line}| => |{y_line}| (target)')
